import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs

class KMeans:
    def __init__(self, k=3, max_iters=100, tol=1e-4, random_state=None):
        """
        KMeans聚类算法
        
        参数:
        k: 聚类数量
        max_iters: 最大迭代次数
        tol: 收敛阈值（质心变化小于该值时停止迭代）
        random_state: 随机种子
        """
        self.k = k
        self.max_iters = max_iters
        self.tol = tol
        self.random_state = random_state
        self.centroids = None
        self.labels = None
        self.inertia_ = None
        
    def _initialize_centroids(self, X):
        """初始化质心 - 使用KMeans++方法"""
        np.random.seed(self.random_state)
        n_samples, n_features = X.shape
        
        # 第一个质心随机选择
        centroids = [X[np.random.randint(n_samples)]]
        
        # 选择剩余的k-1个质心
        for _ in range(1, self.k):
            # 计算每个样本到最近质心的距离
            distances = np.array([min([np.linalg.norm(x - c)**2 for c in centroids]) 
                                for x in X])
            
            # 根据距离的概率分布选择下一个质心
            probabilities = distances / distances.sum()
            cumulative_probs = probabilities.cumsum()
            r = np.random.rand()
            
            for i, p in enumerate(cumulative_probs):
                if r < p:
                    centroids.append(X[i])
                    break
        
        return np.array(centroids)
    
    def _assign_clusters(self, X, centroids):
        """将样本分配到最近的质心"""
        distances = np.linalg.norm(X[:, np.newaxis] - centroids, axis=2)
        return np.argmin(distances, axis=1)
    
    def _update_centroids(self, X, labels):
        """更新质心为每个簇的均值"""
        new_centroids = np.array([X[labels == i].mean(axis=0) 
                                for i in range(self.k)])
        return new_centroids
    
    def _calculate_inertia(self, X, labels, centroids):
        """计算簇内平方和（inertia）"""
        inertia = 0
        for i in range(self.k):
            cluster_points = X[labels == i]
            if len(cluster_points) > 0:
                inertia += np.sum((cluster_points - centroids[i])**2)
        return inertia
    
    def fit(self, X):
        """
        训练KMeans模型
        
        参数:
        X: 输入数据，形状为(n_samples, n_features)
        """
        n_samples, n_features = X.shape
        
        # 初始化质心
        self.centroids = self._initialize_centroids(X)
        
        # 迭代优化
        for iteration in range(self.max_iters):
            # 分配样本到簇
            self.labels = self._assign_clusters(X, self.centroids)
            
            # 更新质心
            new_centroids = self._update_centroids(X, self.labels)
            
            # 检查收敛
            centroid_shift = np.linalg.norm(new_centroids - self.centroids, axis=1).max()
            
            if centroid_shift < self.tol:
                print(f"收敛于第 {iteration + 1} 次迭代")
                break
                
            self.centroids = new_centroids
        
        # 计算最终的簇内平方和
        self.inertia_ = self._calculate_inertia(X, self.labels, self.centroids)
        
        return self
    
    def predict(self, X):
        """预测新样本的簇标签"""
        return self._assign_clusters(X, self.centroids)
    
    def fit_predict(self, X):
        """训练模型并返回预测标签"""
        self.fit(X)
        return self.labels

# 测试和可视化
def test_kmeans():
    # 生成测试数据
    X, y_true = make_blobs(n_samples=300, centers=4, n_features=2, 
                          random_state=42, cluster_std=0.60)
    
    # 创建并训练KMeans模型
    kmeans = KMeans(k=4, random_state=42)
    labels = kmeans.fit_predict(X)
    
    # 可视化结果
    plt.figure(figsize=(15, 5))
    
    # 原始数据
    plt.subplot(1, 3, 1)
    plt.scatter(X[:, 0], X[:, 1], c=y_true, cmap='viridis', s=50, alpha=0.7)
    plt.title('原始数据')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')
    
    # KMeans聚类结果
    plt.subplot(1, 3, 2)
    plt.scatter(X[:, 0], X[:, 1], c=labels, cmap='viridis', s=50, alpha=0.7)
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], 
                c='red', marker='X', s=200, label='质心')
    plt.title('KMeans聚类结果')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')
    plt.legend()
    
    # 对比真实标签和预测标签
    plt.subplot(1, 3, 3)
    plt.scatter(X[:, 0], X[:, 1], c=labels, cmap='viridis', s=50, alpha=0.7)
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], 
                c='red', marker='X', s=200, label='质心')
    
    # 标记分类错误的点
    incorrect = labels != y_true
    plt.scatter(X[incorrect, 0], X[incorrect, 1], 
                facecolors='none', edgecolors='black', s=100, 
                linewidth=2, label='分类错误')
    
    plt.title('分类错误点标记')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')
    plt.legend()
    
    plt.tight_layout()
    plt.show()
    
    # 打印模型信息
    print(f"簇内平方和 (Inertia): {kmeans.inertia_:.4f}")
    print(f"质心位置:\n{kmeans.centroids}")

# 肘部法则确定最佳K值
def elbow_method(X, max_k=10):
    """使用肘部法则确定最佳聚类数量"""
    inertias = []
    k_range = range(1, max_k + 1)
    
    for k in k_range:
        kmeans = KMeans(k=k, random_state=42)
        kmeans.fit(X)
        inertias.append(kmeans.inertia_)
    
    plt.figure(figsize=(10, 6))
    plt.plot(k_range, inertias, 'bo-')
    plt.xlabel('聚类数量 K')
    plt.ylabel('簇内平方和 (Inertia)')
    plt.title('肘部法则 - 确定最佳K值')
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    # 生成测试数据
    X, _ = make_blobs(n_samples=300, centers=4, n_features=2, 
                     random_state=42, cluster_std=0.60)
    
    # 测试KMeans算法
    test_kmeans()
    
    # 使用肘部法则
    elbow_method(X)